package cc.magicjson.easy.batch.strategy;

import cc.magicjson.easy.batch.core.BatchProcessor;
import cc.magicjson.easy.batch.core.BatchStrategy;
import cc.magicjson.easy.batch.model.BatchContext;
import cc.magicjson.easy.batch.model.BatchResult;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;

import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;

/**
 * 并行批处理策略
 */
@Slf4j
public class ParallelBatchStrategy<T> implements BatchStrategy<T> {
    @Override
    public BatchResult<T> execute(BatchContext<T> context) {
        List<T> items = context.getItems();
        List<List<T>> batches = Lists.partition(items,
            context.getBatchSize());

        ExecutorService executor = Executors.newFixedThreadPool(
            context.getThreadCount());
        BatchResult<T> totalResult = new BatchResult<>();

        try {
            List<CompletableFuture<BatchResult<T>>> futures = batches.stream()
                .map(batch -> CompletableFuture.supplyAsync(() -> {
                    BatchContext<T> batchContext = cloneContext(
                        context, batch);
                    return new TransactionalBatchStrategy<T>().execute(batchContext);
                }, executor))
                .collect(Collectors.toList());

            CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
                .thenAccept(v -> futures.forEach(f ->
                    totalResult.merge(f.join())))
                .join();

        } finally {
            executor.shutdown();
        }

        return totalResult;
    }

    private BatchContext<T> cloneContext(
        BatchContext<T> original, List<T> batchItems) {
        return BatchContext.<T>builder()
            .items(batchItems)
            .sqlKey(original.getSqlKey())
            .parameterExtractor(original.getParameterExtractor())
            .useTransaction(original.isUseTransaction())
            .batchSize(original.getBatchSize())
            .allowPartialSuccess(original.isAllowPartialSuccess())
            .progressCallback(original.getProgressCallback())
            .retryAttempts(original.getRetryAttempts())
            .retryDelay(original.getRetryDelay())
            .jdbcTemplate(original.getJdbcTemplate())
            .transactionTemplate(original.getTransactionTemplate())
            .sqlProvider(original.getSqlProvider())
            .build();
    }
}
